import Foundation
import Combine

class MemoryService {
    // 单例模式
    static let shared = MemoryService()
    
    private init() {}
    
    // 获取当前上下文窗口的消息
    func getContextWindow(messages: [ChatMessage], maxTokens: Int = 4000) -> [ChatMessage] {
        // 简单实现：保留最近的消息，直到达到最大token数
        // 实际应用中，应该使用更复杂的算法来计算token数量
        
        var contextMessages: [ChatMessage] = []
        var estimatedTokens = 0
        
        // 从最新消息开始添加
        for message in messages.reversed() {
            // 估算token数量（粗略估计：每4个字符约1个token）
            let messageTokens = message.content.count / 4 + 20 // 20是角色信息等的开销
            
            if estimatedTokens + messageTokens > maxTokens {
                break
            }
            
            contextMessages.insert(message, at: 0) // 保持原始顺序
            estimatedTokens += messageTokens
        }
        
        return contextMessages
    }
    
    // 生成会话摘要
    func generateSummary(messages: [ChatMessage]) -> String {
        // 在实际应用中，这里应该调用AI来生成摘要
        // 这里只是一个简单的实现
        
        let recentMessages = Array(messages.suffix(min(5, messages.count)))
        var topics = Set<String>()
        
        for message in recentMessages {
            let words = message.content.split(separator: " ")
            for word in words where word.count > 5 {
                topics.insert(String(word))
                if topics.count >= 3 {
                    break
                }
            }
            if topics.count >= 3 {
                break
            }
        }
        
        return "讨论了以下主题: \(topics.joined(separator: ", "))"
    }
    
    // 提取重要信息
    func extractKeyInformation(message: ChatMessage) -> [String] {
        // 在实际应用中，这里应该使用AI或NLP技术来提取关键信息
        // 这里只是一个简单的实现
        
        let sentences = message.content.split(separator: "。").map { String($0) }
        return sentences.filter { $0.count > 20 } // 简单地认为较长的句子包含重要信息
    }
    
    // 构建带有记忆的系统提示
    func buildSystemPromptWithMemory(basePrompt: String, keyInformation: [String]) -> String {
        var systemPrompt = basePrompt
        
        if !keyInformation.isEmpty {
            systemPrompt += "\n\n重要背景信息:\n"
            for (index, info) in keyInformation.enumerated() {
                systemPrompt += "\(index + 1). \(info)\n"
            }
        }
        
        return systemPrompt
    }
    
    // 管理长期记忆
    func manageLongTermMemory(session: inout ChatSession) {
        // 如果消息数量超过阈值，提取关键信息并创建摘要
        if session.messages.count > 50 {
            let oldMessages = Array(session.messages.prefix(session.messages.count - 20))
            let keyInfo = oldMessages.flatMap { extractKeyInformation(message: $0) }
            
            // 创建一个系统消息来保存摘要
            if !keyInfo.isEmpty {
                let summaryContent = "之前对话的重要信息:\n" + keyInfo.joined(separator: "\n")
                let summaryMessage = ChatMessage(
                    content: summaryContent,
                    role: .system,
                    timestamp: Date()
                )
                
                // 保留最近的20条消息，并在开头添加摘要
                session.messages = [summaryMessage] + Array(session.messages.suffix(20))
            }
        }
    }
} 